import argparse
import logging
import os

import requests

from db import impl
from utils import conf_log
from conf import REQ_CONNECT_TIMEOUT, DOWNLOAD_TIMEOUT, DOWNLOAD_DIR, FILENAME_FETCHER_PATTERN, DB_CONF


def get_args():
    parser = argparse.ArgumentParser(description='A downloader of IoT firmware downloading some '
                                                 'from the specified firmware db.')
    parser.add_argument('-n', '--download_nums', type=int, default=100,
                        help='The number of records would be downloaded.')
    parser.add_argument('-p', '--parallel_nums', type=int, default=5,
                        help='The number of worker-parallel-nums would be used to download.')
    parser.add_argument('-b', '--brand', type=str, default='Huawei',
                        help='Specify firmwares of the which brand would be downloaded.')
    return parser.parse_args()


def create_download_dir(dir: str):
    if not os.path.exists(dir):
        os.mkdir(dir)


def download_firmware(url: str, dir: str):
    import threading, time
    cur_t = threading.current_thread().name
    logging.info('[thread=%s, url=%s]Downloading...' % (cur_t, url))
    rsp = requests.get(url=url, stream=True, timeout=REQ_CONNECT_TIMEOUT)  # unit sec
    ret = FILENAME_FETCHER_PATTERN.findall(rsp.headers.get('Content-Disposition', 'filename=""'))
    if ret:
        filename = ret[0]
    else:
        logging.info('[thread=%s, url=%s]Cannot find the filename.' % (cur_t, url))
        return
    filename = os.path.join(dir, filename)
    logging.info('[thread=%s, url=%s]Starts to write bites into file.' % (cur_t, url))
    start_time = time.time()
    try:
        with open(filename, 'wb') as f:
            for chunk in rsp.iter_content(10240):   # 10k
                f.write(chunk)
                cur_time = time.time()
                if cur_time - start_time > DOWNLOAD_TIMEOUT:
                    raise TimeoutError('Downloads too slow and stops the download.')
    except TimeoutError as e:
        if os.path.exists(filename):
            os.remove(filename)
        logging.exception('[thread=%s, url=%s]Downloads the file too slow and give up the download.' % (cur_t, url))
    logging.info('[thread=%s, url=%s]Finished to write into the file.' % (cur_t, url))
    logging.info('[thread=%s, url=%s]Downloaded.' % (cur_t, url))


def download_firmwares(urls: list, parallel: int, dir: str):
    from queue import Queue, Empty
    que = Queue(len(urls))
    for url in urls:
        if url.startswith('http://'):
            que.put_nowait(url)
        else:
            logging.info('Has discarded the url %s bcz its prefix not begins with http://' % url)
    import threading
    size = que.qsize()
    logging.info('Total %s firmwares will be downloaded.' % size)

    def run():
        while True:
            try:
                url_ = que.get_nowait()
                if url_:
                    download_firmware(url=url_, dir=dir)
                    s = que.qsize()
                    logging.info("The current progress is %s%%" % (100.0 * (size - s) / size))
            except Empty as ignore:
                break
        logging.info("The thread %s exited" % threading.current_thread().name)

    for t in [threading.Thread(target=run) for _ in range(parallel)]:
        t.start()


def get_branch_id(brand):
    rst = cli.query('id', FROM='brand', WHERE='name=\'%s\'' % brand)
    if rst:
        return rst[0][0]


def get_images(brand_id, n):
    rst = cli.query('id', FROM='image', WHERE='brand_id=%s' % brand_id, LIMIT=n)
    if rst:
        return [image[0] for image in rst]


def get_urls(image_ids=None, n=None):
    if image_ids:
        ret = []
        for iid in image_ids:
            rst = cli.query('url', FROM='product', WHERE='iid=%s' % iid)
            if rst:
                ret.append(rst[0][0])
        return ret
    elif n:
        rst = cli.query('url', FROM='product', LIMIT=n)
        if rst:
            return [product[0] for product in rst]


if __name__ == '__main__':
    conf_log()

    args = get_args()
    download_nums = args.download_nums
    parallel_nums = args.parallel_nums
    brand = args.brand

    cli = impl.PostgreSQLDBClient(conf=DB_CONF)
    cli.open()

    try:
        if brand:
            bid = get_branch_id(brand)
            if not bid:
                raise RuntimeError('There not exists brand named %s' % brand)
            image_ids = get_images(bid, download_nums)
            if not image_ids:
                raise RuntimeError('There not exists any images.')
            logging.info("The size of image_ids is %s" % len(image_ids))
            urls = get_urls(image_ids=image_ids)
        else:
            urls = get_urls(n=download_nums)

        if not urls:
            raise RuntimeError('There not exists any urls.')
    except Exception as e:
        cli.close()
        raise e

    create_download_dir(DOWNLOAD_DIR)
    # download
    download_firmwares(urls=urls, parallel=parallel_nums, dir=DOWNLOAD_DIR)
